# from elasticsearch import Elasticsearch
from elasticsearch import Elasticsearch
from langchain_elasticsearch import ElasticsearchStore
from langchain_openai import OpenAIEmbeddings

from backend.app.core.conf import settings


class ESConnection(object):
    host = None
    client = None
    embeddings = None

    def __init__(self):
        url = "https://{}:{}".format(settings.EL_HOST, settings.EL_PORT)
        self.host = url

    def open(self):
        host = self.host
        ssl_fingerprint = settings.EL_SSL_FINGERPRINT
        user = settings.EL_USER
        password = settings.EL_PASSWORD
        self.client = Elasticsearch(host, basic_auth=(user, password), ssl_assert_fingerprint=ssl_fingerprint, max_retries=10)
        self.embeddings = OpenAIEmbeddings()

    def close(self):
        self.client.close()

    def get_client(self):
        return self.client

    def get_langchain_client(self, index_name: str):
        db = ElasticsearchStore(
            es_connection=self.client,
            embedding=self.embeddings,
            index_name=index_name,
            vector_query_field="embedding",
            query_field="text"
        )
        return db

    def get_embeddings(self):
        return self.embeddings

    def get_index(self, index_name: str):
        """获取指定索引连接"""
        client = self.client
        if not client.indices.exists(index=index_name):
            client.indices.create(index=index_name)


es_conn = ESConnection()
